# -*- coding: utf8 -*-

from elasticsearch_dsl import Search
from elasticsearch import Elasticsearch


class ES(object):
    def __init__(self, hosts=None):
        self.conn = None
        if hosts:
            self.conn = Elasticsearch([hosts], max_retries=0)

    def init_app(self, app):
        hosts = app.config.get("ELASTICSEARCH_HOST", "localhost:9200")
        if hosts:
            self.conn = Elasticsearch([hosts], max_retries=0)

    # def search(self, index, fields, mapping, **kwargs):
    #     """构建查询"""
    #     res = Search(using=self.conn, index=index).source(fields)
    #     for key in kwargs.keys():
    #         res = res.filter("term", **{mapping[key]: kwargs[key]})
    #     res = res.scan()
    #     return [hit.to_dict() for hit in res]

    def search_hits(self, index, query):
        es = self.conn
        resp = es.search(index, body=query)
        resp_hits = resp['hits']['hits']
        return resp_hits
    # 检测index是否存在
    def index_exists(self, index):
        return self.conn.indices.exists(index=index)

    def add_index(self, index, body=None):
        self.conn.indices.create(index=index, ignore=400, body=body)

    def set_settings(self, index, settings):
        self.conn.indices.get_settings(
            index=index, doc_type=index, body=settings)

    def set_mapping(self, index, mapping):
        self.conn.indices.put_mapping(
            index=index, doc_type=index, body=mapping)

    def add_data(self, index, data):
        self.conn.index(index=index, doc_type=index, body=data)

    def del_index(self, index):
        self.conn.indices.delete(index=index)

    def update_by_query(self, index, update_body):
        r = self.conn.update_by_query(index=index, doc_type=index, body=update_body)
        return r


if __name__ == "__main__":
    # result = ES(hosts="10.249.151.23:9200").search(
    #     index="noah_alarm_yunsuo_20200209",
    #     fields=["unique_id"],
    #     mapping={"key": "unique_id"},
    #     key="1577195407132909334",
    # )
    # print(result)

    updateBody = {
        "query": {
            "bool": {
                "must": [
                    {"term": {"_id": "vulnerability--aa6d1668-fd95-4513-b411-a76547477c97"}}
                ],
            }
        },
        "script": {
            "inline": "ctx._source.tag = params.tag",
            "params": {
                "tag": "chaizc"
            },
            "lang": "painless"

        }
    }

    _es = ES(hosts="10.48.93.17:9200")
    result = _es.update_by_query("skygc_spec_entity_vulnerability", updateBody)
